import tensorflow.keras as keras
import tensorflow as tf

def ConvBlock(input_shape,filters,kernel_size,strides=(1,1),padding = None):
    # 设置padding
    if strides ==(2,2):
        padding='valid'
    else:
        padding ='same'
    # 构建
    # 输入
    inputs = keras.Input(shape=input_shape)
    # 卷积层
    conv = keras.layers.Conv2D(filters,kernel_size,strides,padding,kernel_regularizer=keras.regularizers.l2(5e-4))(inputs)
    # BN层
    bn = keras.layers.BatchNormalization()(conv)
    # 激活
    output = keras.layers.LeakyReLU(alpha=0.1)(bn)
    return keras.Model(inputs,output)

def ResBlock(input_shape,filters,blocks):
    # 输入
    inputs = keras.Input(input_shape)
    # pad
    pad = keras.layers.ZeroPadding2D(padding=((1,0),(1,0)))(inputs)
    # CBL
    results = ConvBlock(pad.shape[1:],filters=filters,kernel_size=(3,3),strides=(2,2))(pad)
    # 构建残差单元
    for i in range(blocks):
        # CBL
        results_conv = ConvBlock(results.shape[1:],filters=filters//2,kernel_size=(1,1))(results)
        # CBL
        results_conv = ConvBlock(results_conv.shape[1:],filters=filters,kernel_size=(3,3))(results_conv)
        # add
        results = keras.layers.Add()([results_conv,results])
    return keras.Model(inputs,results)

def Body(input_shape,name):
    # 输入
    inputs = keras.Input(input_shape)
    # CBL
    cb = ConvBlock(inputs.shape[1:],filters=32,kernel_size=(3,3))(inputs)
    # RES1
    rb1 = ResBlock(cb.shape[1:],filters=64,blocks=1)(cb)
    rb2 = ResBlock(rb1.shape[1:],filters=128,blocks=2)(rb1)
    rb3 = ResBlock(rb2.shape[1:],filters=256,blocks=8)(rb2)
    rb4 = ResBlock(rb3.shape[1:],filters=512,blocks=8)(rb3)
    rb5 = ResBlock(rb4.shape[1:],filters=1024,blocks=8)(rb4)
    return keras.Model(inputs,outputs = (rb5,rb4,rb3),name=name)


def Output(input_shape,input_filters,output_filters,name):
    # 输入
    inputs = keras.Input(input_shape)
    # 串联CBL
    cb1 = ConvBlock(inputs.shape[1:],filters=input_filters,kernel_size=(1,1))(inputs)
    cb2 = ConvBlock(cb1.shape[1:],filters=input_filters*2,kernel_size=(3,3))(cb1)
    cb3 = ConvBlock(cb2.shape[1:],filters=input_filters,kernel_size=(1,1))(cb2)
    cb4 = ConvBlock(cb3.shape[1:],filters=input_filters*2,kernel_size=(3,3))(cb3)
    cb5 = ConvBlock(cb4.shape[1:],filters=input_filters,kernel_size=(1,1))(cb4)
    cb6 = ConvBlock(cb5.shape[1:],filters=input_filters*2,kernel_size=(3,3))(cb5)
    cb7 = ConvBlock(cb6.shape[1:],filters=output_filters,kernel_size=(1,1))(cb6)
    return keras.Model(inputs,outputs = (cb5,cb7),name=name)

def YOLOv3(input_shape,classnum=2):
    # anchor数目
    anchor_num = 3
    # 输入数据
    inputs = keras.Input(shape=input_shape,name="input")
    # 通过backbone获取特征图
    large,middle,small =Body(inputs.shape[1:],name='backbone')(inputs)
    # 大目标
    x1,y1 = Output(large.shape[1:],512,anchor_num*(classnum+5),name="output1")(large)
    # 输出
    y1 = keras.layers.Reshape((input_shape[0]//32,input_shape[1]//32,3,5+classnum),name = "reshape1")(y1)
    # 中等目标
    cb1 = ConvBlock(x1.shape[1:],256,kernel_size=(1,1))(x1)
    # 上采样
    us1 = keras.layers.UpSampling2D(2,name='upsampling1')(cb1)
    # 融合
    cat1 = keras.layers.Concatenate()([us1,middle])
    # 输出
    x2,y2 = Output(cat1.shape[1:],256,anchor_num*(classnum+5),name='output2')(cat1)
    y2 = keras.layers.Reshape((input_shape[0]//16,input_shape[1]//16,3,5+classnum),name='reshape2')(y2)
    # 小目标
    cb2 = ConvBlock(x2.shape[1:],128,kernel_size=(1,1))(x2)
    # 上采样
    us2 = keras.layers.UpSampling2D(2,name='upsampling2')(cb2)
    # 融合
    cat2 = keras.layers.Concatenate()([us2,small])
    # 输出
    x3,y3 = Output(cat2.shape[1:],128,anchor_num*(classnum+5),name='output3')(cat2)
    y3 = keras.layers.Reshape((input_shape[0]//8,input_shape[1]//8,3,5+classnum),name='reshape3')(y3)
    return keras.Model(inputs,outputs=(y1,y2,y3))


def OutputParser(input_shape, img_shape, anchors):
    # 输入数据
    feats = tf.keras.Input(input_shape)
    # 对回归结果进行调整
    # 1.计算网格左上角坐标
    grid_y = tf.keras.layers.Lambda(lambda x: tf.tile(tf.reshape(tf.range(tf.cast(tf.shape(
        x)[1], dtype=tf.float32), dtype=tf.float32), (-1, 1, 1, 1)), (1, tf.shape(x)[2], 1, 1)))(feats)
    grid_x = tf.keras.layers.Lambda(lambda x: tf.tile(tf.reshape(tf.range(tf.cast(tf.shape(
        x)[2], dtype=tf.float32), dtype=tf.float32), (1, -1, 1, 1)), (tf.shape(x)[1], 1, 1, 1)))(feats)
    # 构建grid的网格表示
    grid = tf.keras.layers.Concatenate(axis=-1)([grid_x, grid_y])
    # 2.计算box_xy
    box_xy = tf.keras.layers.Lambda(lambda x: (tf.math.sigmoid(x[0][..., 0:2]) + x[1]) / tf.cast(
        [tf.shape(x[1])[1], tf.shape(x[1])[0]], dtype=tf.float32))([feats, grid])
    # 3.计算box_wh
    box_wh = tf.keras.layers.Lambda(lambda x, y, z: tf.math.exp(x[..., 2:4]) * y / tf.cast(
        [z[1], z[0]], dtype=tf.float32), arguments={'y': anchors, 'z': img_shape})(feats)
    # 4.是否包含目标的置信度
    box_confidence = tf.keras.layers.Lambda(
        lambda x: tf.math.sigmoid(x[..., 4]))(feats)
    # 5.分类
    box_class_probs = tf.keras.layers.Lambda(
        lambda x: tf.math.sigmoid(x[..., 5:]))(feats)
    # 返回结果
    return tf.keras.Model(inputs=feats, outputs=(box_xy, box_wh, box_confidence, box_class_probs))


if __name__ == "__main__":
    yolov3 = YOLOv3((416, 416, 3), 80)
    yolov3.summary()